# 文档管理类

from datetime import datetime
from .models import Document, Category
from .vector_store import VectorStore

class DocumentManager:
    """文档管理类，提供文档和分类的管理功能"""
    
    def __init__(self, db):
        """初始化文档管理器"""
        self.db = db
        # 用于更新ChromaDB的VectorStore实例
        self.vector_store = VectorStore(db)
    
    def create_category(self, name, description=None):
        """创建分类"""
        category = Category(name=name, description=description)
        self.db.execute(category, commit=True)
        return category.id
    
    def get_all_categories(self):
        """获取所有分类"""
        session = self.db.connect()
        try:
            categories = session.query(Category).order_by(Category.name).all()
            # 在会话关闭前准备返回数据
            result = [(category.id, category.name, category.description) for category in categories]
            return result
        finally:
            # 确保在访问完所有需要的属性后再关闭会话
            session.close()
    
    def get_category_by_id(self, category_id):
        """通过ID获取分类"""
        session = self.db.connect()
        try:
            category = session.query(Category).filter(Category.id == category_id).first()
            if category:
                # 在会话关闭前返回数据
                return (category.id, category.name, category.description)
            return None
        finally:
            # 确保在访问完所有需要的属性后再关闭会话
            session.close()
    
    def update_category(self, category_id, name, description=None):
        """更新分类信息"""
        session = self.db.connect()
        try:
            category = session.query(Category).filter(Category.id == category_id).first()
            if not category:
                return False
            
            # 更新分类信息
            category.name = name
            category.description = description
            
            session.commit()
            return True
        finally:
            session.close()
    
    def update_document_categories(self, document_id, category_ids):
        """更新文档的分类（支持多个分类）"""
        session = self.db.connect()
        document = session.query(Document).filter(Document.id == document_id).first()
        if document:
            # 清空现有分类
            document.categories.clear()
            
            # 添加新分类
            if category_ids:
                categories = (
                    session.query(Category).filter(Category.id.in_(category_ids)).all()
                )
                document.categories.extend(categories)
            
            document.updated_at = datetime.now()
            session.commit()
            session.close()
            # 同时更新ChromaDB中的元数据
            try:
                # 重用初始化时创建的VectorStore实例
                self.vector_store.collection.update(
                    ids=[str(document_id)],
                    metadatas=[
                        {
                            "category_ids": (
                                ",".join(map(str, category_ids)) if category_ids else ""
                            )
                        }
                    ],
                )
            except Exception as e:
                print(f"更新ChromaDB元数据时出错: {e}")
    
    def get_documents_by_category(self, category_id):
        """获取指定分类下的所有文档"""
        session = self.db.connect()
        try:
            if category_id:
                # 通过多对多关系查询属于指定分类的文档
                category = (
                    session.query(Category).filter(Category.id == category_id).first()
                )
                if not category:
                    return []
                documents = category.documents
            else:
                # 获取所有文档
                documents = session.query(Document).all()
            
            # 返回文档信息，包括所属的所有分类ID
            result = []
            for doc in documents:
                # 确保在会话关闭前访问categories
                category_ids = [cat.id for cat in doc.categories]
                result.append(
                    (
                        doc.id,
                        doc.content,
                        doc.title,
                        category_ids,
                        doc.created_at,
                        doc.updated_at,
                    )
                )
            
            return result
        finally:
            # 确保在访问完所有需要的属性后再关闭会话
            session.close()
    
    def get_document_by_id(self, document_id):
        """通过ID获取文档"""
        session = self.db.connect()
        try:
            document = session.query(Document).filter(Document.id == document_id).first()
            
            if document:
                # 确保在会话关闭前访问categories
                category_ids = [cat.id for cat in document.categories]
                return (
                    document.id,
                    document.content,
                    category_ids,
                    document.title,
                    document.created_at,
                    document.updated_at,
                )
            return None
        finally:
            # 确保在访问完所有需要的属性后再关闭会话
            session.close()
    
    def update_document(self, document_id, title=None, content=None, category_ids=None):
        """更新文档信息"""
        result = False
        session = self.db.connect()
        try:
            document = session.query(Document).filter(Document.id == document_id).first()
            if not document:
                return False
            
            # 更新文档信息
            if title is not None:
                document.title = title
            if content is not None:
                document.content = content
            
            # 更新分类
            if category_ids is not None:
                # 清空现有分类
                document.categories.clear()
                # 添加新分类
                if category_ids:
                    categories = (
                        session.query(Category).filter(Category.id.in_(category_ids)).all()
                    )
                    document.categories.extend(categories)
            
            # 更新时间
            document.updated_at = datetime.now()
            session.commit()
            
            # 在会话关闭前获取需要的数据
            current_categories = [cat.id for cat in document.categories]
            current_title = document.title
            current_content = document.content
            current_id = document.id
            
            result = True
        finally:
            # 确保在访问完所有需要的属性后再关闭会话
            session.close()

        # 更新ChromaDB中的向量和元数据
        if result:
            try:
                # 使用之前保存的数据，不需要再次访问数据库
                # 获取新的嵌入
                embedding = self.vector_store.get_text_embedding(current_content)
                # 更新ChromaDB
                self.vector_store.collection.update(
                    ids=[str(document_id)],
                    embeddings=[embedding],
                    metadatas=[
                        {
                            "document_id": current_id,
                            "category_ids": (
                                ",".join(map(str, current_categories)) if current_categories else ""
                            ),
                            "title": current_title if current_title else "",
                            "content": current_content,
                        }
                    ],
                )
            except Exception as e:
                print(f"更新ChromaDB时出错: {e}")
                
        return result
    
    def delete_document(self, document_id):
        """删除文档"""
        session = self.db.connect()
        document = session.query(Document).filter(Document.id == document_id).first()
        if not document:
            return False
        
        # 从ChromaDB中删除向量
        try:
            self.vector_store.collection.delete(ids=[str(document_id)])
        except Exception as e:
            print(f"从ChromaDB删除向量时出错: {e}")
        
        # 从数据库中删除文档
        session.delete(document)
        session.commit()
        session.close()
        return True
        
    def get_total_documents(self):
        """获取文档总数"""
        session = self.db.connect()
        try:
            # 直接使用count()方法获取文档总数
            total_count = session.query(Document).count()
            return total_count
        finally:
            session.close()